import zmq
import json
import requests
import base64

import threading
import datetime  # 新增导入 datetime 模块

# Doris 的用户名和密码
username = 'root'
password = 'root'
# 认证信息
auth_string = f'{username}:{password}'
auth_bytes = auth_string.encode('utf-8')
encoded_auth = base64.b64encode(auth_bytes).decode('utf-8')
def process_message(data):
    # 根据 type 字段选择表名
    table_name = get_table_name(data['type'])
    fields = get_fields(data['type'])
    # 解析 xdrLog 字段
    xdr_logs = data['message'][0]['xdrLog'].split('\r\n')  # 修改：按 \r\n 分割多条话单
    
    # 构建插入数据
    data_to_insert_list = []
    for log in xdr_logs:
        log_fields = log.split('|')
        # 修改：将 start_time 和 end_time 从微秒转换为日期格式
        if 'start_time' in fields and 'end_time' in fields:
            start_time_index = fields.index('start_time')
            end_time_index = fields.index('end_time')
            log_fields[start_time_index] = datetime.datetime.fromtimestamp(int(log_fields[start_time_index]) / 1000000).strftime('%Y-%m-%d %H:%M:%S.%f')
            log_fields[end_time_index] = datetime.datetime.fromtimestamp(int(log_fields[end_time_index]) / 1000000).strftime('%Y-%m-%d %H:%M:%S.%f')
        
        data_to_insert = ','.join([f'"{log_fields[i]}"' for i in range(len(fields))])
        data_to_insert_list.append(data_to_insert)
    
    # 使用 Stream Load 将数据写入 Doris
    url = f"http://10.10.4.41:18030/api/one_piece/{table_name}/_stream_load"
    headers = {
        "Expect": "100-continue",
        "Content-Type": "text/plain; charset=UTF-8",
        "label": f"label_{data['sequence']}",
        'Authorization': f'Basic {encoded_auth}'
    }
    data_to_insert = '\n'.join(data_to_insert_list)
    # 发送 Stream Load 请求
    response = requests.put(
        url,
        headers=headers,
        data=data_to_insert,
        auth=(username, password)
    )

    # 检查响应状态
    if response.status_code == 200:
        result = response.json()
        if result.get('Status') == 'Success':
            print('数据导入成功')
        else:
            print(f'数据导入失败: {result.get("Message")}')
    else:
        print(f'请求失败，状态码: {response.status_code}, 响应内容: {response.text}')

def listen_socket():
    # 创建一个 ZeroMQ 上下文
    context = zmq.Context()
    # 创建一个 REP（应答）套接字
    socket = context.socket(zmq.REP)
    # 绑定到指定的 TCP 地址和端口
    socket.bind("tcp://*:5555")

    print("正在监听端口 5555...")
    try:
        while True:
            # 接收消息
            message = socket.recv_string()
            print("=============" + message)
            data = json.loads(message)
            
            # 启动一个新的线程来处理消息
            threading.Thread(target=process_message, args=(data,)).start()
            
            socket.send_string("OK, 200")

    except KeyboardInterrupt:
        print("程序已停止")
    finally:
        # 关闭套接字
        socket.close()
        # 终止上下文
        context.term()

def get_table_name(type_value):
    # 根据 type 字段返回表名
    table_mapping = {
        4040007: "td_realtime_record",
        # 添加其他 type 和表名的映射
    }
    return table_mapping.get(type_value, "default_table")

def get_fields(type_value):
    # 根据 type 字段返回字段列表
    fields_mapping = {
        4040007: ['start_time','end_time','ip_type','user_ipv4','user_ipv6','user_port','l4_protocol','app_server_ip_ipv4','app_server_ip_ipv6','app_server_port','app_type','app_subtype','app_content','protocol_type','host','in_netid','out_netid','terminal_name','terminal_type','user_account','duration','ul_ip_packet','ul_data','dl_ip_packet','dl_data','ul_retrans_ip_packet','dl_retrans_ip_packet','tcp12_latency','tcp23_latency','response_time','ul_avg_rtt','dw_avg_rtt','avg_count']
    }
    return fields_mapping.get(type_value, [])

if __name__ == "__main__":
    listen_socket()